import os
import argparse
import re
from azure.reader import AzureReader
import open3d as o3d
import sys
import json
pwd = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(pwd,"Open3D", "examples", "python", "reconstruction_system"))
from initialize_config import initialize_config
 
if __name__ == "__main__":
    # get argument of input date
    parser = argparse.ArgumentParser()
    parser.add_argument('--date', type=str)
    args = parser.parse_args()
    date = args.date
    if date is None:
        raise ValueError('Please input date')
    # /share/hhodataset/azure4/20231030/000/azure4.mkv
    # /share/hhodataset/VTS/20231002/002/VTS_data.npz
    # all files in /share/hhodataset/azure2 + date
    vts_path = os.path.join('/share/hhodataset/VTS', date)

    for folder_name in os.listdir(vts_path):
        try:
            if not re.match(r"\d{3}$", folder_name):
                continue
            vts_data_path = os.path.join(vts_path, folder_name, 'VTS_data.npz')
            azure2_path = os.path.join('/share/hhodataset/azure2', date, folder_name)
            azure3_path = os.path.join('/share/hhodataset/azure3', date, folder_name)
            azure4_path = os.path.join('/share/hhodataset/azure4', date, folder_name)
            azure5_path = os.path.join('/share/hhodataset/azure5', date, folder_name)
            azure_paths = ["0","1",azure2_path, azure3_path, azure4_path, azure5_path]
            if not os.path.exists(vts_data_path):
                print('{} :VTS data not exist '.format(vts_data_path))
                continue
            if "azure2.mkv" not in os.listdir(azure2_path) and "azure2.txt" not in os.listdir(azure2_path):
                print('{} :azure2.mkv not exist '.format(azure2_path))
                continue
            if "azure3.mkv" not in os.listdir(azure3_path) and "azure3.txt" not in os.listdir(azure3_path):
                print('{} :azure3.mkv not exist '.format(azure3_path))
                continue
            if "azure4.mkv" not in os.listdir(azure4_path) and "azure4.txt" not in os.listdir(azure4_path):
                print('{} :azure4.mkv not exist '.format(azure4_path))
                continue
            if "azure5.mkv" not in os.listdir(azure5_path) and "azure5.txt" not in os.listdir(azure5_path):
                print('{} :azure5.mkv not exist '.format(azure5_path))
                continue
            print('start process {}'.format(folder_name))
            
            for i in range(2,6):
                # reader = AzureReader({
                #     "output":os.path.join(azure_paths[i], "vis"),
                #     "input": os.path.join(azure_paths[i], "azure{}.mkv".format(i)),
                #     "convert_mkv": True,
                # })
                # reader.start()
                output = os.path.join(azure_paths[i], "vis")
                input = os.path.join(azure_paths[i], "azure{}.mkv".format(i))
                if not os.path.exists(output):
                    os.mkdir(output)
                    os.mkdir('{}/color'.format(output))
                    os.mkdir('{}/depth'.format(output))
                reader = o3d.io.AzureKinectMKVReader()
                reader.open(input)
                metadata = reader.get_metadata()
                o3d.io.write_azure_kinect_mkv_metadata(
                    '{}/intrinsic.json'.format(output), metadata)
                config = {
                    'path_dataset': output,
                    'path_intrinsic': '{}/intrinsic.json'.format(output)
                }
                initialize_config(config)
                with open('{}/config.json'.format(output), 'w') as f:
                    json.dump(config, f, indent=4)
                    
                idx = 0
                while not reader.is_eof():
                    rgbd = reader.next_frame()
                    if rgbd is None:
                        continue
                    # if not vis_geometry_added:
                    #     vis.add_geometry(rgbd)
                    #     vis_geometry_added = True

                    if output is not None:
                        color_filename = '{0}/color/{1:05d}.jpg'.format(
                            output, idx)
                        print('Writing to {}'.format(color_filename))
                        o3d.io.write_image(color_filename, rgbd.color)

                        depth_filename = '{0}/depth/{1:05d}.png'.format(
                            output, idx)
                        print('Writing to {}'.format(depth_filename))
                        o3d.io.write_image(depth_filename, rgbd.depth)
                        idx += 1
                reader.close()
        except Exception as e:
            print("Wrong at {}. {}".format(folder_name, str(e)))
            continue
        
        
        